from autogen import ConversableAgent, GroupChat, GroupChatManager, UserProxyAgent, AssistantAgent
from utils.configuration import load_config, set_environment, default, exists
config = load_config('./env.yaml')

# llm_config = {"config_list": [{"model": 'gpt-3.5-turbo',
#                                "api_key": 'sk-ypDOrct5sWOZEVCUYWmNqK8rsGXcJ6svajXmNrPdEeU5tshv',
#                                'base_url': 'https://api.chatanywhere.tech' #config['llm']['base_url']
#                                }]}
llm_config = {"config_list": [{"model": config['llm']['model'],
                               "api_key": 'sk-ypDOrct5sWOZEVCUYWmNqK8rsGXcJ6svajXmNrPdEeU5tshv',
                               'base_url': config['llm']['base_url']
                               }]}

userAgent = UserProxyAgent(
    name='admin',
    llm_config=llm_config,
    code_execution_config=False,
    human_input_mode="NEVER"
)

planner = AssistantAgent(
    name="planner",
    system_message="""你将扮演一个助手
                   功能描述：根据给出的章节，请记住目前写到第几章了，请输出下一章应该撰写哪章的内容。
                   你的输出格式 "你需要分小节的章节为:[章节名称]"，如果你认为写完了请输出【结束写作】。非常重要的一点，请不要输出任何额外的东西
                   不需要你描述章节内容，你只需要给出章节标题。
                   """,
    description="planner：对论文的章节进行安排，包括已完成章节、未完成章节、下一章",
    llm_config=llm_config,
    human_input_mode="NEVER"
)
sectionAgent = AssistantAgent(
    name='SectionAgent',
    system_message="""
    你将扮演一个写作框架助手，
    输入格式：输入格式有两种情况，1. 你需要分小节的章节为:[章节名称]； 2. 【撰写完成小节X】
    功能描述：当输入是1. 你需要分小节的章节为:[章节名称]；你需要根据输入的大标题进行小节设计，我会根据你给我的小节进行写作，当输入是2. 【撰写完成小节X】，请你记住我写到了哪一个小节并且告诉我下一小节要写什么。
    你的输出: "【章节名称：XXX：该章节有如下小节：[小节1，小节2，小节3....]，请撰写小节：X】" 当所有的小节都撰写完毕时请输出"【请问下一章需要撰写什么？】"
    如果请不要输出任何额外的东西！！！不需要你描述小节内容，你只需要给出章节标题。 """,
    llm_config=llm_config,
    human_input_mode="NEVER",
    description="SectionAgent：会对给出将章节分成若干小节，并且指导contentAgent进行小节撰写"
)

contentAgent = AssistantAgent(
    name='ContentAgent',
    system_message="""
    你将扮演一个写作assistant，
    功能描述：你需要根据我给你的小节标题，请使用markdown格式进行写作。内容要非常非常非常的详细，包括对小节标题的详细介绍，小节内容撰写。
    举个例子比如小节标题是“系统功能需求分析”，首先你需要描述什么是系统功能需求分析，其次你对系统进行详细的功能需求分析。
    输入格式：【章节名称：XXX：该章节有如下小节：[小节1，小节2，小节3....]，请撰写小节：XXXX】
    你的输出：“【小节正文内容】 此外你还要输出 是否完成这一小节的撰写 若完成了撰写请输出【撰写完成小节X】如果还没有请输出 【未完待续】 ” 
    """,
    llm_config=llm_config,
    description="ContentAgent：根据给出的章节小节进行写作",
    human_input_mode="NEVER"
)

transition = {
    userAgent: [planner],
    planner: [sectionAgent],
    sectionAgent: [contentAgent,planner],
    contentAgent: [ sectionAgent]
}

groupchat = GroupChat(
    agents=[sectionAgent, contentAgent, planner, userAgent],
    messages=[],
    max_round=40,
    send_introductions=True,
    speaker_transitions_type="allowed",
    allowed_or_disallowed_speaker_transitions=transition,


)
manager = GroupChatManager(groupchat=groupchat, llm_config=llm_config,
                           system_message="你是一个调度的各个专家的assistant，你需要调度ContentAgent，planner，sectionAgent进行论文写作"
                                          "planner会对论文章节进行安排，并且记住目前的进度，和下一章的安排，"
                                          "sectionAgent会根据planner的给出的下一章安排进行小节的安排"
                                          "contentAgent会对每个小节进行撰写，请你合理调度这三者帮助我完成论文撰写.",
                           is_termination_msg= lambda msg : "结束写作" in msg['content']
                           )

if __name__ == '__main__':
    # chat_history = sectionAgent.initiate_chat(
    #     contentAgent,
    #     message=""
    #
    # )
    chat_history = userAgent.initiate_chat(
        manager,
        message="帮我写一篇，基于Vus.js的美食分享管理系统的开发与实践。"
                "请从以下一下六章进行撰写。未完成章节【1.绪论；2.技术分析；3.需求分析；4.功能实现；5.系统演示；6.总结；】",
    )
